import argparse  # 解析命令行参数
import os        # 文件路径操作
import numpy as np
import torch
from transformers import AutoTokenizer, LlamaForCausalLM  # 加载 LLaMA 模型与分词器


from lib.prune import run_pruning
from lib.eval import evaluate_ppl_light  # 轻量版 PPL 评估（适合剪枝后快速测试）
# from lib.eval import evaluate_ppl      # 若需要更完整的评估可以启用


def parse_args():
    parser = argparse.ArgumentParser(description='Prune and evaluate LLaMA models')

    parser.add_argument('--model', type=str, required=True, help='LLaMA模型的路径（HF格式）')
    parser.add_argument('--cache_dir', type=str, default='llm_weights', help='模型缓存文件夹路径')
    parser.add_argument('--fp16', action='store_true', help='是否使用 float16 精度加载模型')

    parser.add_argument('--method', choices=['flap', 'wanda_sp', 'mag_sp'], default='flap', help='剪枝方法选择')
    parser.add_argument('--pruning_ratio', type=float, default=0.0, help='全局剪枝比例')
    parser.add_argument('--remove_heads', type=int, default=0, help='删除多少个注意力头')
    parser.add_argument('--metric', choices=['IFV', 'WIFV', 'WIFN', 'N/A'], default='WIFV', help='FLAP使用的评分指标')
    parser.add_argument('--structure', choices=['UL-UM', 'UL-MM', 'AL-MM', 'AL-AM', 'N/A'], default='AL-AM', help='FLAP的结构类型')
    parser.add_argument('--nsamples', type=int, default=512, help='采样数据的数量（用于剪枝统计）')
    parser.add_argument('--seed', type=int, default=0, help='随机种子，确保实验可复现')
    parser.add_argument('--unstructured', action='store_true', help='是否启用非结构化剪枝')

    parser.add_argument('--eval', action='store_true', help='是否进行剪枝后的PPL评估')
    parser.add_argument('--save_dir', type=str, default=None, help='剪枝后模型保存目录')

    parser.add_argument('--test_file', type=str, default='/home/mht/FLAP/dataset/wikitext-raw/wiki.test.raw', help='测试文本路径')
    return parser.parse_args()


def set_seed(seed: int):
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


def setup_model(path: str, cache_dir: str, fp16: bool):
    # 选择精度
    dtype = torch.float16 if fp16 else torch.float32

    # 从 HuggingFace 格式加载 LLaMA 模型（low_cpu_mem_usage 避免内存爆炸）
    model = LlamaForCausalLM.from_pretrained(
        path,
        torch_dtype=dtype,
        cache_dir=cache_dir,
        low_cpu_mem_usage=True,
        device_map='auto' if torch.cuda.is_available() else None  # 自动选择GPU或CPU
    )
    model.eval()  # 设置为评估模式，关闭 dropout 等
    return model


def main():
    args = parse_args()          # 解析命令行参数
    set_seed(args.seed)          # 设置随机种子，确保剪枝一致

    print(f"[INFO] Loading model from {args.model}, fp16={args.fp16}")
    model = setup_model(args.model, args.cache_dir, args.fp16)
    tokenizer = AutoTokenizer.from_pretrained(args.model, use_fast=False)

    device = next(model.parameters()).device  # 获取模型所在设备
    print(f"[INFO] Using device: {device}")

    print("[INFO] Running pruning...")
    run_pruning(args, model, tokenizer)  # 调用 lib/prune.py 中的剪枝主函数
    print("[INFO] Pruning finished.")

    total = sum(p.numel() for p in model.parameters())              # 所有参数数量
    nonzero = sum(torch.count_nonzero(p).item() for p in model.parameters())  # 非零参数数量
    sparsity = 1 - nonzero / total                                  # 计算稀疏率
    print(f"[RESULT] Sparsity: {sparsity:.4f}, Total params: {total/1e9:.2f}B")

    if args.eval:
        ppl = evaluate_ppl_light(model, tokenizer, device, test_file=args.test_file, max_tokens=1000)
        print(f"[RESULT] Lightweight PPL: {ppl:.2f}")


    if args.save_dir:
        try:
            os.makedirs(args.save_dir, exist_ok=True)
            model.save_pretrained(args.save_dir)       # 保存模型权重
            tokenizer.save_pretrained(args.save_dir)   # 保存分词器配置
            print(f"[INFO] Saved pruned model to {args.save_dir}")
        except Exception as e:
            print(f"[ERROR] Failed to save model: {e}")

if __name__ == '__main__':
    main()
